
from datetime import datetime
import pandas as pd

import akshare as ak
import backtrader as bt

def get_etf_hist_data(symbol, start_date="20240701", end_date="20240702"):
    """
    获取 etf 历史数据，并转换为 backtrader 需要的数据格式
    """

    etf_df = ak.fund_etf_hist_min_em(symbol=symbol, start_date=start_date, end_date=end_date)
    
    #etf_df.to_csv()

    # 将交易日期设置为索引值
    etf_df.index = pd.to_datetime(etf_df["时间"])
    etf_df.sort_index(inplace=True)
    etf_df.drop(axis=1, columns='时间', inplace=True)

    recon_data = {'High': etf_df.最高, 'Low': etf_df.最低, 'Open': etf_df.开盘, 'Close': etf_df.收盘, \
                  'Volume': etf_df.成交量}
    df_recon = pd.DataFrame(recon_data)

    print(df_recon)

    return df_recon


class CSimpleDualMAStrategy(bt.Strategy):
    """
    简单的双均线交易策略
    """

    # 全局设定交易策略的参数
    params=(
            ('ma_short',30),
            ('ma_long', 60),
           )
    
    def __init__(self):

        # 指定价格序列
        self.dataclose=self.datas[0].close
        # 初始化交易指令、买卖价格和手续费
        self.order = None
        self.buyprice = None
        self.buycomm = None

        # 添加移动均线指标
        # 5日移动平均线
        self.sma5 = bt.indicators.SimpleMovingAverage(self.datas[0], period=self.params.ma_short)
        # 10日移动平均线
        self.sma10 = bt.indicators.SimpleMovingAverage(self.datas[0], period=self.params.ma_long)

        self.rsi5 = bt.indicators.RSI_EMA(self.datas[0], period=self.params.ma_short,subplot=True)
        self.william_ad = bt.indicators.WilliamsAD(self.datas[0], subplot=True)

    def log(self, txt, dt=None, doprint=False):
        # 日志函数，用于统一输出日志格式
        if doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print('%s, %s' % (dt.isoformat(), txt))

    def notify_order(self, order):
        """
        用于跟踪交易指令（order）的状态。order具有提交，接受，买入/卖出执行和价格，已取消/拒绝等状态。

        订单状态处理
        Arguments:
            order {object} -- 订单状态
        """
        if order.status in [order.Submitted, order.Accepted]:
            # 如订单已被处理，则不用做任何事情
            return
        # 检查订单是否完成
        if order.status in [order.Completed]:
            if order.isbuy():
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
                self.log(f'买入:\n价格:{order.executed.price},\
                                成本:{order.executed.value},\
                                手续费:{order.executed.comm}', doprint=True)
            else:
                self.log(f'卖出:\n价格：{order.executed.price},\
                                成本: {order.executed.value},\
                                手续费{order.executed.comm}', doprint=True)
            self.bar_executed = len(self)
        # 订单因为缺少资金之类的原因被拒绝执行
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')
        # 订单状态处理完成，设为空
        self.order = None

    # 记录交易收益情况
    def notify_trade(self, trade):
        """
        交易成果
        Arguments:
            trade {object} -- 交易状态
        """
        if not trade.isclosed:
            return
        # 显示交易的毛利率和净利润
        self.log('Operation Profit, Gross%.2f, Net %.2f' %
                 (trade.pnl, trade.pnlcomm), doprint=True)

    def next(self):
        # 记录收盘价
        self.log('Close, %.2f' % self.dataclose[0])
        # Access -1, because drawdown[0] will be calculated after "next"
        #self.log('DrawDown: %.2f' % self.stats.drawdown.drawdown[-1], doprint=True)
        #self.log('MaxDrawDown: %.2f' % self.stats.drawdown.maxdrawdown[-1], doprint=True)

        if self.order: # 是否有指令等待执行
            return
        # 是否持仓
        if not self.position: # 没有持仓
            # 执行买入条件判断：MA5上扬突破MA10，买入
            if self.sma5[0] > self.sma10[0]:
                self.order = self.buy() # 执行买入
                self.log('Buy Create, %.2f' % self.dataclose[0], doprint = True)
        else:
            # 执行卖出条件判断：MA5下穿跌破MA10，卖出
            if self.sma5[0] < self.sma10[0]:
                self.order = self.sell() # 执行卖出
                self.log('Sell Create, %.2f' % self.dataclose[0], doprint = True)

    # 回测结束后输出
    def stop(self):
        self.log(u'金叉死叉策略 %2d Ending Value %.2f' %
                 (self.params.ma_long, self.broker.getvalue()), doprint=True)
        

if __name__ == "__main__":
    """
    设置Cerebro引擎：关于策略回测，把数据和策略添加到Cerebro中之外，还有设置一些参数。
    比如broker的设置，像初始资金、交易佣金，也可以用addsizer设定每次交易买入的股数。
    """

    start_date = datetime(2024, 6, 1)
    end_date = datetime(2024, 7, 2)

    ## 获取沪深300ETF（510300）数据
    etf_df = get_etf_hist_data(symbol="510300", start_date=start_date.strftime("%Y%m%d"), end_date=end_date.strftime("%Y%m%d"))
    data = bt.feeds.PandasData(dataname=etf_df, fromdate=start_date,todate=end_date)
    cerebro = bt.Cerebro() # 初始化cerebro回测系统设置 默认 stdstats=True
    cerebro.adddata(data) # 将数据传入回测系统
    cerebro.addstrategy(CSimpleDualMAStrategy) # 将交易策略加载到回测系统中
    cerebro.broker.setcash(10000) # 设置初始资本为10,000
    cerebro.addsizer(bt.sizers.FixedSize, stake=500) # 设定每次交易买入的股数
    cerebro.broker.setcommission(commission=0.002) # 设置交易手续费为 0.2%
    cerebro.addobserver(bt.observers.DrawDown)
    print('初始资金: %.2f' % cerebro.broker.getvalue())

    results = cerebro.run() #运行回测系统
    cerebro.plot()